import torch.nn as nn
import torch


class FCB(nn.Module):
    def __init__(self, cnum):
        super(FCB, self).__init__()
        self.cnum = cnum

    def _delete(self, vector, index, n):

        v_left = vector[:index]
        v_right = vector[index + n:]
        v = torch.cat((v_left, v_right))

        return v

    def forward(self, features, target):
        loss_sum = 0

        batch_size = features.size(0)
        if float(batch_size) % 2 != 0:
            raise Exception('Incorrect batch size provided')
        #     batch_size = math.floor(batch_size % 2)*2 #如果batchsize不能被2整除，则舍弃一个

        target_left = target[:int(0.5 * batch_size)]
        target_right = target[int(0.5 * batch_size):]

        batch_left = features[:int(0.5 * batch_size)]
        batch_right = features[int(0.5 * batch_size):]

        for i in range(batch_size // 2):
            # ground-truth
            index_left = target_left[i]
            index_right = target_right[i]

            if index_left == index_right:

                left = self._delete(batch_left[i], index_left, self.cnum)
                right = self._delete(batch_right[i], index_left, self.cnum)
            else:

                left = self._delete(batch_left[i], index_left, self.cnum)
                left = self._delete(left, index_right, self.cnum)

                right = self._delete(batch_left[i], index_left, self.cnum)
                right = self._delete(right, index_left, self.cnum)

            loss = torch.norm((left - right).abs(), 2, 0)

            loss_sum = loss_sum + loss
            out = loss_sum / torch.tensor(batch_size)
        return out
